from fastapi import FastAPI, BackgroundTasks
from fastapi.responses import StreamingResponse
import uuid
import asyncio
import os
from openai import AsyncOpenAI


os.environ["OPENAI_API_KEY"] = "myAPI_KEY"  
client = AsyncOpenAI()

app = FastAPI()


FUNCTIONS = {
    "zh2en": "中文 -> 英文翻译",
    "en2zh": "英文 -> 中文翻译",
    "summarize": "文本总结"
}

TASK_RESULTS = {}


async def call_llm(prompt: str) -> str:

    response = await client.chat.completions.create(
        model="gpt-4o-mini",  
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content.strip()



@app.get("/functions")
async def get_functions():
    return {"functions": FUNCTIONS}



@app.post("/invoke/{func_name}")
async def invoke_function(func_name: str, text: str):
    if func_name not in FUNCTIONS:
        return {"error": "未知功能"}
    if func_name == "zh2en":
        prompt = f"Translate the following Chinese into English:\n{text}"
    elif func_name == "en2zh":
        prompt = f"Translate the following English into Chinese:\n{text}"
    elif func_name == "summarize":
        prompt = f"Summarize the following text:\n{text}"
    else:
        return {"error": "未实现功能"}

    result = await call_llm(prompt)
    return {"result": result}



@app.post("/task/{func_name}")
async def submit_task(func_name: str, text: str, background_tasks: BackgroundTasks):
    if func_name not in FUNCTIONS:
        return {"error": "未知功能"}
    
    task_id = str(uuid.uuid4())
    TASK_RESULTS[task_id] = None

    async def run_task():
        if func_name == "zh2en":
            prompt = f"Translate the following Chinese into English:\n{text}"
        elif func_name == "en2zh":
            prompt = f"Translate the following English into Chinese:\n{text}"
        elif func_name == "summarize":
            prompt = f"Summarize the following text:\n{text}"
        else:
            TASK_RESULTS[task_id] = "未实现功能"
            return

        result = await call_llm(prompt)
        TASK_RESULTS[task_id] = result

    background_tasks.add_task(run_task)
    return {"task_id": task_id}



@app.get("/task/result/{task_id}")
async def get_task_result(task_id: str):
    if task_id not in TASK_RESULTS:
        return {"error": "任务不存在"}
    result = TASK_RESULTS[task_id]
    if result is None:
        return {"status": "pending"}
    return {"status": "done", "result": result}



@app.post("/stream/{func_name}")
async def stream_function(func_name: str, text: str):
    if func_name not in FUNCTIONS:
        return {"error": "未知功能"}

    if func_name == "zh2en":
        prompt = f"Translate the following Chinese into English:\n{text}"
    elif func_name == "en2zh":
        prompt = f"Translate the following English into Chinese:\n{text}"
    elif func_name == "summarize":
        prompt = f"Summarize the following text:\n{text}"
    else:
        return {"error": "未实现功能"}

    async def event_generator():
        stream = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            stream=True,  
        )
        async for chunk in stream:
            delta = chunk.choices[0].delta.content
            if delta:
                yield delta

    return StreamingResponse(event_generator(), media_type="text/plain")
